import asyncio
import typing
from asyncio import ALL_COMPLETED, FIRST_COMPLETED, FIRST_EXCEPTION


class AsyncContext:
    def __init__(self, return_when=ALL_COMPLETED):
        self.return_when = return_when

    async def __aenter__(self):
        self.task_list: typing.List[asyncio.futures.Future] = []
        return self

    def run(self, func, callback=None, args: tuple = None, kwargs: dict = None):
        if args is None:
            args = tuple()
        if kwargs is None:
            kwargs = dict()

        task = asyncio.create_task(func(*args, **kwargs))
        if callback:
            task.add_done_callback(callback)
        self.task_list.append(task)

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.task_list:
            await asyncio.wait(self.task_list, return_when=self.return_when)
